package com.classactionpl.jaas.ldap;

import java.io.IOException;
import java.security.Principal;
import java.security.acl.Group;
import java.util.HashSet;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

import javax.naming.Context;
import javax.naming.NamingEnumeration;
import javax.naming.NamingException;
import javax.naming.directory.Attribute;
import javax.naming.directory.Attributes;
import javax.naming.directory.DirContext;
import javax.naming.directory.InitialDirContext;
import javax.naming.directory.SearchControls;
import javax.naming.directory.SearchResult;
import javax.security.auth.Subject;
import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.PasswordCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.auth.login.FailedLoginException;
import javax.security.auth.login.LoginException;
import javax.security.auth.spi.LoginModule;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import com.classactionpl.jaas.UserGroup;
import com.classactionpl.jaas.UserPrincipal;

/**
 * A BindingLoginModule uses the JNDI directory services to bind a principal for
 * the purposes of authentication. If a successful bind is achieved then the
 * groups of the directory are searched for instances where the principal is a
 * member. Groups where the principal is a member are returned as roles.
 * 
 * Particular focus has been placed on performance. Only one bind is required
 * and that is for the user being authenticated. One lookup on the bound context
 * and then one subtree search to determine member groups is all that is
 * performed. Furthermore only attributes that are required are returned from
 * lookups and searches.
 * 
 * The underlying assumption with this module is that directory access is locked
 * down and that the principal can bind to their own entry. In addition the
 * principals may view the attribute that resolves to a username. It is also
 * expected that all bound users are able to view the naming attribute of all
 * group entries that can be potentially returned as roles.
 * 
 * For example, ensure that the base DN of where the users bind to is an
 * administrative area:
 * 
 * <pre>
 * administrativeRole: accessControlSpecificArea
 * </pre>
 * 
 * and that it has an ACI similar to the following:
 * 
 * <pre>
 * {
 *     identificationTag &quot;enableSearchForThisUser&quot;,
 *     precedence 14,
 *     authenticationLevel simple,
 *     itemOrUserFirst userFirst: 
 *     {
 *         userClasses { thisEntry },
 *         userPermissions 
 *         {
 *             {
 *                 precedence 14,
 *                 protectedItems { allUserAttributeTypesAndValues, entry },
 *                 grantsAndDenials 
 *                 {
 *                     grantReturnDN,
 *                     grantBrowse,
 *                     grantRead 
 *                 }
 *             }
 *         }
 *     }
 * }
 * </pre>
 * 
 * The above allows the user that has bound to this entry to browse and read its
 * attributes - but no one elses.
 * 
 * Similarly for where the base of the groups are you require and
 * administrativeRole and an ACI that allows all bound users to browse and read
 * attributes (you could get specific about which attributes):
 * 
 * <pre>
 * {
 *     identificationTag &quot;enableSearchForAllUsers&quot;,
 *     precedence 14,
 *     authenticationLevel simple,
 *     itemOrUserFirst userFirst: 
 *     {
 *         userClasses { allUsers },
 *         userPermissions 
 *         {
 *             {
 *                 precedence 14,
 *                 protectedItems { allUserAttributeTypesAndValues, entry },
 *                 grantsAndDenials 
 *                 {
 *                     grantReturnDN,
 *                     grantBrowse,
 *                     grantRead 
 *                 }
 *             }
 *         }
 *     }
 * }
 * </pre>
 * 
 * <p>
 * Here is a list of JAAS parameters that are required:
 * 
 * <table>
 * <tr>
 * <td>initialContextFactory</td>
 * <td>The fully distinguished JNDI context factory class name</td>
 * </tr>
 * <tr>
 * <td>providerURL</td>
 * <td>The URL required to connect to the directory</td>
 * </tr>
 * <tr>
 * <td>securityAuthentication</td>
 * <td>"none", "simple" or <i>sasl_mech</i></td>
 * </tr>
 * <tr>
 * <td>userBase</td>
 * <td>The base DN of the directory to perform one-level binding of users.</td>
 * </tr>
 * <tr>
 * <td>userRDN</td>
 * <td>The attribute used to determine the RDN of the user.</td>
 * </tr>
 * <tr>
 * <td>userReturnAttributeAsUserName</td>
 * <td>The name of the attribute returned from the bind's associated lookup.
 * This used to identify the username.</td>
 * </tr>
 * <tr>
 * <td>groupSearchBase</td>
 * <td>The part of the directory to perform subtree searches on for groups
 * potentially with an association to a user.</td>
 * </tr>
 * <tr>
 * <td>groupSearchFilter</td>
 * <td>The directory search specification. {0} will be substituted with the
 * bound user's DN.</td>
 * </tr>
 * <tr>
 * <td>groupSearchReturnAttributeAsGroupName</td>
 * <td>The name of the attribute returned from the search that is used for the
 * name of the group.</td>
 * </tr>
 * </table>
 * 
 * A sample JAAS configuration file might look like:
 * 
 * <pre>
 * ldaploginmodule {
 *  com.classactionpl.jaas.ldap.BindingLoginModule required
 *  debug=&quot;true&quot;
 * 	initialContextFactory=&quot;com.sun.jndi.ldap.LdapCtxFactory&quot;
 * 	providerURL=&quot;ldap://localhost:10389&quot;
 * 	securityAuthentication=&quot;simple&quot;
 * 	userBase=&quot;ou=users,ou=system&quot;
 * 	userRDN=&quot;uid&quot;
 * 	userReturnAttributeAsUserName=&quot;uid&quot;
 * 	groupSearchBase=&quot;ou=groups,ou=system&quot;
 * 	groupSearchFilter=&quot;(member={0})&quot;
 * 	groupSearchReturnAttributeAsGroupName=&quot;cn&quot;
 * ;};
 * </pre>
 * 
 * @author huntc
 * 
 */
public class BindingLoginModule implements LoginModule {

	Log logger = LogFactory.getLog(BindingLoginModule.class);

	Hashtable<Object, Object> env;

	String userBase;
	String userRDN;
	String userReturnAttributeAsUserName;

	String groupSearchBase;
	String groupSearchFilter;
	String groupSearchReturnAttributeAsGroupName;

	Subject subject;
	CallbackHandler callbackHandler;

	Principal principal;
	Set<Principal> groups;

	boolean loginSucceeded = false;
	boolean commitSucceeded = false;

	public boolean abort() throws LoginException {
		boolean abortSucceeded = loginSucceeded;

		logout();

		return abortSucceeded;
	}

	/**
	 * This method is responsible for the ldap interaction. In essence the
	 * user's username is formed into a DN and then a bind is attempted. If it
	 * succeeds then the attribute describing the user's username is queried
	 * for.
	 * 
	 * Authentication is deemed successful if the above completes. The second
	 * part, and therefore an optional part, is to determine the groups that the
	 * user's bound DN belongs to. A subtree search is used for this.
	 * 
	 * @param principalName
	 * @param credentials
	 * @return
	 * @throws LoginException
	 */
	@SuppressWarnings("unchecked")
	boolean authenticate(String principalName, byte[] credentials)
			throws LoginException {
		boolean authenticated = false;

		// Bind to our ldap repo. If we succeed then authentication is
		// successful.

		String userDN = userRDN + "=" + principalName + "," + userBase;

		Hashtable<Object, Object> localEnv = (Hashtable<Object, Object>) env
				.clone();
		localEnv.put(Context.SECURITY_PRINCIPAL, userDN);
		localEnv.put(Context.SECURITY_CREDENTIALS, credentials);
		DirContext ctx;
		try {
			ctx = new InitialDirContext(localEnv);
			try {
				Attributes attrs = ctx.getAttributes(userDN,
						new String[] { userReturnAttributeAsUserName });
				if (attrs != null) {

					NamingEnumeration<? extends Attribute> attrsIter = attrs
							.getAll();
					if (attrsIter.hasMore()) {
						Attribute userNameAttr = attrsIter.next();
						if (userNameAttr != null) {

							principal = new UserPrincipal((String) userNameAttr
									.get());

							if (logger.isDebugEnabled()) {
								logger.debug("Found user: "
										+ principal.getName());
							}

							authenticated = true;

							// Got the user, now, while we're still bound as the
							// user, let's search the groups to find out who the
							// user belongs to. In a way, finding groups is a
							// bonus and is not a requirement of the module that
							// has to be fulfilled i.e. you can authenticate and
							// not belong to any groups. (not that this might be
							// terribly practical).

							groups = new HashSet<Principal>();

							SearchControls ctls = new SearchControls();
							ctls.setSearchScope(SearchControls.SUBTREE_SCOPE);
							ctls
									.setReturningAttributes(new String[] { groupSearchReturnAttributeAsGroupName });
							NamingEnumeration<SearchResult> results = ctx
									.search(groupSearchBase, groupSearchFilter,
											new Object[] { userDN }, ctls);

							while (results.hasMore()) {
								SearchResult result = results.next();

								attrs = result.getAttributes();
								if (attrs != null) {

									attrsIter = attrs.getAll();
									if (attrsIter.hasMore()) {
										Attribute groupNameAttr = attrsIter
												.next();
										if (groupNameAttr != null) {
											Group group = new UserGroup(
													(String) groupNameAttr
															.get());
											group.addMember(principal);
											groups.add(group);

											if (logger.isDebugEnabled()) {
												logger.debug("Found group: "
														+ group.getName());
											}
										}
									}
								}
							}
						}
					}
				}

			} finally {
				ctx.close();
			}

		} catch (NamingException e) {
			throw new LoginException(e.toString());
		}

		return authenticated;
	}

	public boolean commit() throws LoginException {
		assert commitSucceeded == false;

		if (loginSucceeded) {
			Set<Principal> existingPrincipals = subject.getPrincipals();
			existingPrincipals.add(principal);
			existingPrincipals.addAll(groups);

			commitSucceeded = true;
		}

		return commitSucceeded;
	}

	public void initialize(Subject subject, CallbackHandler callbackHandler,
			Map<String, ?> sharedState, Map<String, ?> options) {
		// Build our jndi environment options from the options passed in
		String initialContextFactory = (String) options
				.get("initialContextFactory");
		String providerURL = (String) options.get("providerURL");
		String securityAuthentication = (String) options
				.get("securityAuthentication");
		userBase = (String) options.get("userBase");
		userRDN = (String) options.get("userRDN");
		userReturnAttributeAsUserName = (String) options
				.get("userReturnAttributeAsUserName");
		groupSearchBase = (String) options.get("groupSearchBase");
		groupSearchFilter = (String) options.get("groupSearchFilter");
		groupSearchReturnAttributeAsGroupName = (String) options
				.get("groupSearchReturnAttributeAsGroupName");

		env = new Hashtable<Object, Object>();
		env.put(Context.INITIAL_CONTEXT_FACTORY, initialContextFactory);
		env.put(Context.PROVIDER_URL, providerURL);
		env.put(Context.SECURITY_AUTHENTICATION, securityAuthentication);

		this.subject = subject;
		this.callbackHandler = callbackHandler;
	}

	public boolean login() throws LoginException {
		assert loginSucceeded == false;

		Callback callbacks[] = new Callback[2];
		NameCallback nameCallback = new NameCallback("Username: ");
		callbacks[0] = nameCallback;
		PasswordCallback passwordCallback = new PasswordCallback("Password: ",
				false);
		callbacks[1] = passwordCallback;

		try {
			callbackHandler.handle(callbacks);

			String password = new String(passwordCallback.getPassword());
			passwordCallback.clearPassword();

			loginSucceeded = authenticate(nameCallback.getName(), password
					.getBytes("US-ASCII"));

			if (!loginSucceeded) {
				throw new FailedLoginException();
			}

		} catch (IOException e) {
			throw new LoginException(e.toString());
		} catch (UnsupportedCallbackException e) {
			throw new LoginException(e.toString());
		}

		return loginSucceeded;
	}

	public boolean logout() throws LoginException {
		if (commitSucceeded) {
			Set<Principal> existingPrincipals = subject.getPrincipals();
			Iterator<Principal> existingPrincipalsIter = existingPrincipals
					.iterator();
			while (existingPrincipalsIter.hasNext()) {
				Principal existingPrincipal = existingPrincipalsIter.next();
				if (existingPrincipal == principal) {
					existingPrincipalsIter.remove();
				} else {
					Iterator<Principal> groupsIter = groups.iterator();
					while (groupsIter.hasNext()) {
						Principal group = groupsIter.next();
						if (existingPrincipal == group) {
							existingPrincipalsIter.remove();
							groupsIter.remove();
							break;
						}
					}
				}
			}
		}

		loginSucceeded = false;
		commitSucceeded = false;

		principal = null;
		groups = null;

		return true;
	}

}
